#include "crypto_common.h"

// 生成加密上下文
int generate_encrypt_context(crypto_context_t *ctx, const char *password) {
    secure_random_init();
    
    // 设置复杂度级别
    ctx->complexity_level = 3; // 高复杂度
    
    // 生成随机盐值
    generate_secure_random(ctx->salt, SALT_LENGTH);
    
    // 生成随机种子
    ctx->seed = (uint32_t)time(NULL) ^ (uint32_t)rand();
    
    // 如果提供了密码，使用密码生成主密钥
    if (password && strlen(password) > 0) {
        unsigned char hash_output[32];
        hash_function((const unsigned char*)password, strlen(password), hash_output);
        
        // 扩展哈希到完整密钥长度
        for (int i = 0; i < KEY_LENGTH; i++) {
            ctx->primary_key[i] = hash_output[i % 32] ^ (unsigned char)(i * 37);
        }
    } else {
        // 生成随机主密钥
        generate_secure_random(ctx->primary_key, KEY_LENGTH);
    }
    
    // 派生所有轮密钥
    derive_keys(ctx);
    
    return 1;
}

// 保存加密上下文到文件
int save_encrypt_context(const char *filename, crypto_context_t *ctx) {
    FILE *fp = fopen(filename, "wb");
    if (!fp) {
        printf("无法创建密钥文件: %s\n", filename);
        return 0;
    }
    
    // 写入魔术字节标识
    const char magic[] = "SECKEY01";
    fwrite(magic, 1, 8, fp);
    
    // 写入上下文数据
    fwrite(ctx, sizeof(crypto_context_t), 1, fp);
    
    fclose(fp);
    return 1;
}

// 加密文件
int encrypt_file(const char *input_file, const char *output_file, crypto_context_t *ctx) {
    FILE *fp_in = fopen(input_file, "rb");
    if (!fp_in) {
        printf("无法打开输入文件: %s\n", input_file);
        return 0;
    }
    
    FILE *fp_out = fopen(output_file, "wb");
    if (!fp_out) {
        printf("无法创建输出文件: %s\n", output_file);
        fclose(fp_in);
        return 0;
    }
    
    // 写入加密文件头
    const char header[] = "ENCRYPTED_FILE_V2";
    fwrite(header, 1, strlen(header), fp_out);
    
    // 写入盐值（解密时需要）
    fwrite(ctx->salt, 1, SALT_LENGTH, fp_out);
    fwrite(&ctx->seed, sizeof(uint32_t), 1, fp_out);
    fwrite(&ctx->complexity_level, sizeof(uint8_t), 1, fp_out);
    
    // 获取文件大小
    fseek(fp_in, 0, SEEK_END);
    long file_size = ftell(fp_in);
    fseek(fp_in, 0, SEEK_SET);
    
    // 写入原始文件大小
    fwrite(&file_size, sizeof(long), 1, fp_out);
    
    unsigned char buffer[BUFFER_SIZE];
    size_t bytes_read;
    long processed = 0;
    
    printf("开始加密文件...\n");
    
    while ((bytes_read = fread(buffer, 1, BUFFER_SIZE, fp_in)) > 0) {
        // 分块加密
        for (size_t i = 0; i < bytes_read; i += BLOCK_SIZE) {
            size_t block_size = (i + BLOCK_SIZE <= bytes_read) ? BLOCK_SIZE : (bytes_read - i);
            advanced_encrypt_block(buffer + i, block_size, ctx);
        }
        
        fwrite(buffer, 1, bytes_read, fp_out);
        processed += bytes_read;
        
        // 显示进度
        printf("\r进度: %.1f%%", (double)processed / file_size * 100);
        fflush(stdout);
    }
    
    printf("\n加密完成！\n");
    
    fclose(fp_in);
    fclose(fp_out);
    return 1;
}

int main() {
    crypto_context_t ctx;
    char input_file[512];
    char output_file[512];
    char key_file[512];
    char password[256];
    int use_password;
    
    printf("=== 高级文件加密工具 ===\n\n");
    
    printf("是否使用密码? (1=是, 0=否): ");
    scanf("%d", &use_password);
    
    if (use_password) {
        printf("请输入密码: ");
        scanf("%s", password);
    }
    
    printf("请输入要加密的文件路径: ");
    scanf("%s", input_file);
    
    printf("请输入加密后的文件路径: ");
    scanf("%s", output_file);
    
    printf("请输入密钥文件路径: ");
    scanf("%s", key_file);
    
    // 生成加密上下文
    if (!generate_encrypt_context(&ctx, use_password ? password : NULL)) {
        printf("生成加密上下文失败！\n");
        return 1;
    }
    
    // 加密文件
    if (!encrypt_file(input_file, output_file, &ctx)) {
        printf("文件加密失败！\n");
        return 1;
    }
    
    // 保存密钥文件
    if (!save_encrypt_context(key_file, &ctx)) {
        printf("保存密钥文件失败！\n");
        return 1;
    }
    
    printf("\n加密成功完成！\n");
    printf("加密文件: %s\n", output_file);
    printf("密钥文件: %s\n", key_file);
    printf("请妥善保管密钥文件，丢失将无法解密！\n");
    
    // 清除内存中的敏感数据
    memset(&ctx, 0, sizeof(crypto_context_t));
    memset(password, 0, sizeof(password));
    
    return 0;
}